const std = @import("std");
const mem = std.mem;
const math = std.math;
const lbmath = @import("./math/lbmath.zig");
inline fn getOrZero(comptime T: type, A: *const SparseMatrix(T), i: u32, j: u32) T {
    if (A.get(i, j)) |p| return p.*; // 解引用
    return 0;
}
/// 存储LU分解的结果: P*A = L*U
pub fn LUFactors(comptime T: type) type {
    return struct {
        /// 在一个稀疏矩阵中紧凑存储 L (下三角，对角线为1) 和 U (上三角)
        lu: SparseMatrix(T),
        /// 行置换向量 P。p[i] = k 表示计算后的第 i 行是原矩阵的第 k 行。
        p: []u32,
        allocator: mem.Allocator,
        const Self = @This();
        pub fn deinit(self: *Self) void {
            self.lu.deinit();
            self.allocator.free(self.p);
        }

        /// 使用LU分解的结果求解线性方程组 Ax = b, 结果存入 x
        /// 求解过程分为两步:
        /// 1. 解 Ly = P*b (前向替换)
        /// 2. 解 Ux = y (反向替换)
        pub fn solve(self: *const Self, b: []const T, x: []T) !void {
            const n = self.lu.rows;
            std.debug.assert(b.len == n and x.len == n);

            // 1. 应用置换 P 到 b, 结果存入 x (x 此时作为临时向量 y)
            for (0..n) |i| {
                x[i] = b[self.p[i]];
            }

            // 2. 前向替换: 解 Ly = Pb, L的对角线为1
            for (0..n) |i| {
                var sum: T = 0;
                // 遍历 L 的第 i 行的非零元素（不包括隐式的对角线1）
                for (self.lu.data[i].items) |entry| {
                    if (entry.i < i) { // entry.i 是列索引
                        sum += entry.v * x[entry.i];
                    }
                }
                x[i] -= sum;
            }

            // 3. 反向替换: 解 Ux = y
            var i: i64 = n - 1;
            while (i >= 0) : (i -= 1) {
                var sum: T = 0;
                var diag_val: ?T = null;

                // 遍历 U 的第 i 行的非零元素
                for (self.lu.data[@intCast(i)].items) |entry| {
                    if (entry.i > i) {
                        sum += entry.v * x[entry.i];
                    } else if (entry.i == i) {
                        diag_val = entry.v;
                    }
                }

                const pivot = diag_val orelse return error.SingularMatrix;
                if (lbmath.approxEqAbs(T, pivot, 0, null)) {
                    return error.SingularMatrix;
                }
                x[@intCast(i)] = (x[@intCast(i)] - sum) / pivot;
            }
        }
    };
}
pub fn SparseMatrix(comptime T: type) type {
    return struct {
        const SparseMatrixEntry = struct {
            i: u32, // 列索引
            v: T, // 值
        };
        allocator: mem.Allocator,
        rows: u32,
        cols: u32,
        data: []std.ArrayList(SparseMatrixEntry),
        const Self = @This();

        pub fn init(al: mem.Allocator, rows: u32, cols: u32) !Self {
            var data_list = try al.alloc(std.ArrayList(SparseMatrixEntry), rows);
            errdefer al.free(data_list);

            for (0..rows) |i| {
                data_list[i] = .empty;
            }

            return Self{
                .allocator = al,
                .rows = rows,
                .cols = cols,
                .data = data_list,
            };
        }
        /// 释放稀疏矩阵占用的所有内存
        pub fn deinit(self: *Self) void {
            for (0..self.rows) |i| {
                self.data[i].deinit(self.allocator);
            }
            self.allocator.free(self.data);
            self.rows = 0;
            self.cols = 0;
        }
        /// 初始化一个新的、空的稀疏矩阵
        /// 调整矩阵大小，会清空所有数据
        pub fn resize(self: *Self, new_rows: u32, new_cols: u32) !void {
            if (self.rows == new_rows and self.cols == new_cols) return;
            self.deinit();
            self.rows = new_rows;
            self.cols = new_cols;
            self.data = try self.allocator.alloc(std.ArrayList(SparseMatrixEntry), new_rows);
            errdefer self.allocator.free(self.data);

            for (0..new_rows) |i| {
                self.data[i] = .empty;
            }
        }
        /// 创建一个当前矩阵的深拷贝
        pub fn clone(self: *const Self) !Self {
            var new_matrix = try Self.init(self.allocator, self.rows, self.cols);
            errdefer new_matrix.deinit();

            for (self.data, 0..) |row, i| {
                try new_matrix.data[i].appendSlice(self.allocator, row.items);
            }
            return new_matrix;
        }
        /// 将另一个矩阵的数据深拷贝到当前矩阵
        pub fn copyFrom(self: *Self, other: *const Self) !void {
            try self.resize(other.rows, other.cols);
            for (other.data, 0..) |row, i| {
                try self.data[i].appendSlice(self.allocator, row.items);
            }
        }
        /// 矩阵转置，返回一个新的转置矩阵
        pub fn transpose(self: *const Self) !Self {
            var transposed = try Self.init(self.allocator, self.cols, self.rows);
            errdefer transposed.deinit();

            // 填充转置矩阵
            for (self.data, 0..) |row, i| {
                for (row.items) |entry| {
                    try transposed.data[entry.i].append(self.allocator, .{ .i = @intCast(i), .v = entry.v });
                }
            }
            return transposed;
        }

        ///矩阵乘以一个标量
        pub fn scale(self: *Self, factor: T) void {
            for (self.data) |*row| {
                for (row.items) |*entry| {
                    entry.v *= factor;
                }
            }
        }

        /// 矩阵乘以一个向量: result = self * vec
        pub fn multiplyVec(self: *const Self, vec: []const T, result: []T) void {
            std.debug.assert(self.cols <= vec.len); // Can be larger if vec is part of a larger buffer
            std.debug.assert(self.rows == result.len);
            @memset(result, 0);
            for (self.data, 0..) |row, i| {
                for (row.items) |*entry| {
                    result[i] += entry.v * vec[entry.i];
                }
            }
        }

        /// 矩阵加法: self = self + other
        pub fn addEqual(self: *Self, other: *const Self) !void {
            std.debug.assert(self.rows == other.rows and self.cols == other.cols);

            // *** 修正: 使用 alloc 和 free 来管理临时密集向量 ***
            const temp_row = try self.allocator.alloc(T, self.cols);
            defer self.allocator.free(temp_row);

            for (0..self.rows) |i| {
                // 1. 将两行的值合并到一个临时密集行中
                @memset(temp_row, 0.0);
                for (self.data[i].items) |*entry| {
                    temp_row[entry.i] += entry.v;
                }

                for (other.data[i].items) |*entry| {
                    temp_row[entry.i] += entry.v;
                }

                // 2. 从密集行重建稀疏行
                self.data[i].clearRetainingCapacity();
                // clearRetainingCapacity();
                for (temp_row, 0..) |v, j| {
                    if (!lbmath.approxEqAbs(T, v, 0, null)) {
                        try self.data[i].append(self.allocator, .{ .i = @intCast(j), .v = v });
                    }
                }
            }
        }

        /// 矩阵减法: self = self - other
        pub fn subEqual(self: *Self, other: *const Self) !void {
            std.debug.assert(self.rows == other.rows and self.cols == other.cols);

            // *** 修正: 使用 alloc 和 free 来管理临时密集向量 ***
            var temp_row = try self.allocator.alloc(T, self.cols);
            defer self.allocator.free(temp_row);

            for (0..self.rows) |i| {
                @memset(temp_row, 0.0);
                for (self.data[i].items) |entry| {
                    temp_row[entry.i] += entry.v;
                }
                for (other.data[i].items) |entry| {
                    temp_row[entry.i] -= entry.v;
                }

                self.data[i].clearRetainingCapacity();
                for (temp_row, 0..) |v, j| {
                    if (!lbmath.approxEqAbs(T, v, 0, null)) {
                        try self.data[i].append(self.allocator, .{ .i = @intCast(j), .v = v });
                    }
                }
            }
        }

        /// 矩阵乘法: result = m1 * m2
        pub fn multiply(result: *Self, m1: *const SparseMatrix(T), m2: *const SparseMatrix(T)) !void {
            std.debug.assert(m1.cols == m2.rows);
            try result.resize(m1.rows, m2.cols);

            // *** 修正: 使用 alloc 和 free 来管理临时密集向量 ***
            var temp_row = try result.allocator.alloc(T, m2.cols);
            defer result.allocator.free(temp_row);

            for (0..m1.rows) |i| {
                @memset(temp_row, 0.0);

                // 计算结果的第 i 行 (密集形式)
                for (m1.data[i].items) |entry1| {
                    const k = entry1.i;
                    const v1 = entry1.v;
                    for (m2.data[k].items) |entry2| {
                        temp_row[entry2.i] += v1 * entry2.v;
                    }
                }

                // 将密集行转换为稀疏行并存入结果
                result.data[i].clearRetainingCapacity();
                for (temp_row, 0..) |v, j| {
                    if (!lbmath.approxEqAbs(T, v, 0, null)) {
                        try result.data[i].append(result.allocator, .{ .i = @intCast(j), .v = v });
                    }
                }
            }
        }

        /// 获取矩阵元素 (row, col) 的值。对于稀疏矩阵效率不高，仅在必要时使用。
        pub fn get(self: *const Self, row: u32, col: u32) ?*T {
            if (row >= self.rows) return null;
            // 假设行内元素按列索引排序
            for (self.data[row].items) |*entry| {
                if (entry.i == col) return &(entry.v);
                if (entry.i > col) break; // 优化：已排序，提前退出
            }
            return null; // 如果未找到，代表值为0
        }

        /// 设置矩阵元素 (row, col) 的值。
        /// 如果值为0，会尝试移除该元素。
        pub fn set(self: *Self, row: u32, col: u32, value: T) !void {
            std.debug.assert(row < self.rows and col < self.cols);
            const row_list = &self.data[row];
            // 寻找插入点或现有元素
            var idx: usize = 0;
            while (idx < row_list.items.len and row_list.items[idx].i < col) {
                idx += 1;
            }
            const is_zero = lbmath.approxEqAbs(T, value, 0, null);
            if (idx < row_list.items.len and row_list.items[idx].i == col) {
                // 找到了现有元素
                if (is_zero) {
                    _ = row_list.orderedRemove(idx); // 如果新值为0，则删除
                } else {
                    row_list.items[idx].v = value; // 否则更新
                }
            } else {
                // 未找到，如果值不为0，则插入新元素
                if (!is_zero) {
                    try row_list.insert(self.allocator, idx, .{ .i = col, .v = value });
                }
            }
        }

        /// 交换矩阵的两行。这是一个高效的操作，只交换内部的 ArrayList 结构体。
        pub fn swapRows(self: *Self, r1: u32, r2: u32) void {
            if (r1 == r2) return;
            std.debug.assert(r1 < self.rows and r2 < self.rows);
            const tmp = self.data[r1];
            self.data[r1] = self.data[r2];
            self.data[r2] = tmp;
        }
        /// 对稀疏矩阵进行 LU 分解（部分主元选择/行交换）
        /// 返回包含 L、U 与 P 的 LUFactors
        pub fn luDecomposition(self: *const SparseMatrix(T)) !LUFactors(T) {
            std.debug.assert(self.rows == self.cols);
            const n = self.rows;

            var lu_factors = LUFactors(T){
                .lu = try self.clone(),
                .p = try self.allocator.alloc(u32, n),
                .allocator = self.allocator,
            };
            errdefer lu_factors.deinit();

            for (0..n) |i| lu_factors.p[i] = @intCast(i);

            // k 轮消元
            for (0..n) |k| {
                // a) 列 k 上从 k..n-1 寻找绝对值最大的主元
                var max_val: T = @as(T, 0);
                var pivot_row: u32 = k;
                for (k..n) |i| {
                    const a_ik = getOrZero(T, &lu_factors.lu, @intCast(i), k);
                    const absval = @abs(a_ik);
                    if (absval > max_val) {
                        max_val = absval;
                        pivot_row = @intCast(i);
                    }
                }
                if (lbmath.approxEqAbs(T, max_val, @as(T, 0), null)) return error.SingularMatrix;

                // b) 行交换
                if (pivot_row != k) {
                    lu_factors.lu.swapRows(k, pivot_row);
                    const tmp = lu_factors.p[k];
                    lu_factors.p[k] = lu_factors.p[pivot_row];
                    lu_factors.p[pivot_row] = tmp;
                }

                // 现在主元在 (k,k)
                const pivot_ptr = lu_factors.lu.get(k, k) orelse return error.SingularMatrix;
                const pivot_val: T = pivot_ptr.*;

                if (lbmath.approxEqAbs(T, pivot_val, @as(T, 0), null)) return error.SingularMatrix;

                // c) 计算 L 之第 k 列（严格下三角）
                for (k + 1..n) |i| {
                    const a_ik = getOrZero(T, &lu_factors.lu, @intCast(i), k);
                    if (!lbmath.approxEqAbs(T, a_ik, @as(T, 0), null)) {
                        try lu_factors.lu.set(@intCast(i), k, a_ik / pivot_val);
                    }
                }

                // d) 更新右下角子块（Schur complement），填充值可能在这里产生
                // 使用第 k 行的 U (列 > k 的项) 与第 k 列的 L (行 > k 的项)
                for (k + 1..n) |i| {
                    const l_ik = getOrZero(T, &lu_factors.lu, @intCast(i), k);
                    if (lbmath.approxEqAbs(T, l_ik, @as(T, 0), null)) continue;

                    // 遍历当前第 k 行的非零项；只取 U 部分（列 > k）
                    for (lu_factors.lu.data[k].items) |ukj_entry| {
                        const j = ukj_entry.i;
                        if (j <= k) continue; // 只看上三角
                        const u_kj = ukj_entry.v;

                        const old = getOrZero(T, &lu_factors.lu, @intCast(i), j);
                        const new = old - l_ik * u_kj;
                        // 若接近 0，set 会把 0 删掉（你的 set() 正好有这个逻辑）
                        try lu_factors.lu.set(@intCast(i), j, new);
                    }
                }
            }

            return lu_factors;
        }
    };
}

/// 共轭梯度法求解线性方程组 Ax = b
/// A 必须是对称正定矩阵
/// x 是初始猜测，同时也是最终解的输出位置
pub fn computeConjugateGradient(
    comptime T: type,
    allocator: mem.Allocator,
    A: *const SparseMatrix(T),
    b: []const T,
    x: []T,
    max_iter: u32,
    epsilon: T,
) !void {
    std.debug.assert(A.rows == A.cols);
    std.debug.assert(A.rows == b.len and A.rows == x.len);

    const n = A.rows;

    // 分配临时向量
    var r = try allocator.alloc(T, n);
    defer allocator.free(r);
    const p = try allocator.alloc(T, n);
    defer allocator.free(p);
    const Ap = try allocator.alloc(T, n);
    defer allocator.free(Ap);

    // r = b - A*x (计算初始残差)
    A.multiplyVec(x, Ap); // 使用 Ap 作为临时存储
    for (b, Ap, 0..) |b_i, Ax_i, i| {
        r[i] = b_i - Ax_i;
    }

    // p = r
    // *** 修正: 使用 @memcpy 替代 std.mem.copy ***
    // @memcpy 是语言内建函数，接口稳定且性能最佳。
    // 它需要指针和以字节为单位的长度。
    @memcpy(p, r);

    var rs_old = lbmath.inner_product(r, r);

    if (rs_old < epsilon * epsilon) {
        return;
    }

    for (0..max_iter) |_| {
        // Ap = A * p
        A.multiplyVec(p, Ap);

        const alpha_num = rs_old;
        const alpha_den = lbmath.inner_product(p, Ap);
        if (lbmath.approxEqAbs(T, alpha_den, 0, null)) break; // 避免除以零
        const alpha = alpha_num / alpha_den;

        // x = x + alpha * p
        for (x, p) |*x_i, p_i| {
            x_i.* += alpha * p_i;
        }

        // r = r - alpha * Ap
        for (r, Ap) |*r_i, Ap_i| {
            r_i.* -= alpha * Ap_i;
        }

        const rs_new = lbmath.inner_product(r, r);

        // 检查收敛性
        if (rs_new < epsilon * epsilon) {
            break;
        }

        const beta = rs_new / rs_old;
        rs_old = rs_new;

        // p = r + beta * p
        for (p, r) |*p_i, r_i| {
            p_i.* = r_i + beta * p_i.*;
        }
    }
}
